Skip to content

Preserve rows when concatenating pandas ArrayXD batches - #8867

Open
pavanmudumba wants to merge 2 commits into
huggingface:mainfrom
pavanmudumba:codex/fix-arrayxd-concatenation
Open

pavanmudumba wants to merge 2 commits into
huggingface:mainfrom
pavanmudumba:codex/fix-arrayxd-concatenation

Conversation

@pavanmudumba

Copy link
Copy Markdown
Contributor

PandasArrayExtensionArray._concat_same_type can turn whole input batches into individual rows. For example, concatenating dense batches with 2 and 1 rows returns 2 rows instead of 3, with an extra batch dimension inside each row. A single input chunk is also wrapped as one row. Equally sized object-backed chunks take the vstack branch and gain an extra dimension.

Reproduction:

import numpy as np
from datasets.features.features import PandasArrayExtensionArray

a = PandasArrayExtensionArray(np.arange(8, dtype="int32").reshape(2, 2, 2))
b = PandasArrayExtensionArray(np.arange(8, 12, dtype="int32").reshape(1, 2, 2))
result = PandasArrayExtensionArray._concat_same_type([a, b])
print(len(result), result._data.shape)
# Before: 2 (2,)
# After:  3 (3, 2, 2)

Compare each chunk's per-row shape rather than its full batch shape, and concatenate compatible backing arrays along axis 0. For mixed dense/object-backed chunks, allocate one object slot per input row and copy rows individually. This preserves ragged arrays and null rows without retaining the batch grouping.

Added 24 regression/control cases covering unequal and equal batch lengths, a single chunk, empty chunks, int32/float32, Array2D/Array3D, ragged arrays, mixed dense/object storage, null rows, and conversion back to Arrow. Tests exercise the extension-array concatenation hook directly; the separate pandas dtype-comparison issue is tracked in #8375 and #8464.

Validation on Windows / Python 3.12:

  • With the original upstream concatenation method substituted in memory: 16 regression cases fail and 8 controls pass.
  • python -m pytest tests/features/test_array_xd.py tests/test_formatting.py -q -rs: 155 passed, 29 skipped for optional dependencies or platform requirements. A short --basetemp path avoids Windows cache-lock filename limits.
  • Ruff lint and format checks pass for both changed files.
  • git diff --check passes.

@pavanmudumba

Copy link
Copy Markdown
Contributor Author

Hi! Following up on this fix when you have a chance to review. It preserves row counts and shapes when concatenating pandas ArrayXD batches, with regression coverage for dense, ragged, empty, and null-containing inputs. The reproduction and local test results are in the description. Happy to adjust the implementation or add any further coverage you’d find useful. Thanks for your time!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant